import os
import torch
import argparse
import torch.nn.functional as F
from torch.utils.data import DataLoader
from src.models.modnet import MODNet
from src.utils import filter_parameters_by_type
from src.trainer import supervised_training_iter, soc_adaptation_iter
from dataset.data import DatasetNoLabel, DatasetMatting
from tensorboardX import SummaryWriter
import logging
logging.getLogger().setLevel(logging.INFO)

# 输入参数
parser = argparse.ArgumentParser(
    'SOC Training'
)

parser.add_argument('--log', 
    default='./logs/tb_kaggle_kernel_size_04',
    help='tensorboard log path'
)

parser.add_argument('--save_folder', 
    default='./ckpt_kaggle_kernel_size_04/',
    help='checkpoint path'
)

parser.add_argument('--data_path', 
    default='/home/chengk/chk/data/removebg-align/normal3',
    help='image data folder'
)

parser.add_argument('--anno_path',
    default='/home/chengk/chk/data/removebg-align/normal3_removebg',
    help='anno folder'
)

parser.add_argument('--log_image_every_step',
    default=500,
    help='as param name shows'
)

parser.add_argument('--lr',
    default=1e-5,
    help='learning rate'
)

parser.add_argument('--pretrained_model', 
    # default='./modnet_webcam_portrait_matting.ckpt',
    default='./ckpt_kaggle/epoch_55.pth',
    help='pretrained model path'
)

args = parser.parse_args()

tb_logger = SummaryWriter(args.log)
folder_ckpt = args.save_folder
if not os.path.exists(folder_ckpt):
    os.makedirs(folder_ckpt)

bs = 8         # batch size
lr = args.lr       # learn rate
epochs = 20     # total epochs
log_image_every_step = args.log_image_every_step

device = torch.device('cuda:0')
ckpt_pretrain = args.pretrained_model
ckpt = torch.load(ckpt_pretrain, map_location='cpu')

modnet = torch.nn.DataParallel(MODNet(backbone_pretrained=False)).cuda()
modnet_freeze = torch.nn.DataParallel(MODNet(backbone_pretrained=False)).cuda()
modnet.load_state_dict(ckpt)
modnet.module.freeze_norm()
modnet_freeze.load_state_dict(ckpt)
modnet_freeze.eval()
# 停止BatchNorm的更新
# 停止detail和fusion分支
# train_params = filter_parameters_by_type(modnet, torch.nn.modules.batchnorm.BatchNorm2d)
train_params = modnet.module.parameters()
optimizer = torch.optim.Adam(train_params, lr=lr, betas=(0.9, 0.99))
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(0.25 * epochs), gamma=0.1)

# dataset = DatasetNoLabel(root=args.data_path)
dataset = DatasetMatting(root=args.data_path, root_anno=args.anno_path)
dataloader = DataLoader(dataset, bs, num_workers=2, drop_last=True, shuffle=True)     # NOTE: please finish this function

step_scalar, step_image = 0, 0
for epoch in range(0, epochs):
    for idx, data in enumerate(dataloader):
        log_step = (idx+1) % log_image_every_step == 0
        img, matte, trimap = data
        
        # print('overlap_anchor:', overlap_anchor, '\toverlap_target:', overlap_target)
        img = img.to(device)
        matte = matte.to(device)
        trimap = trimap.to(device)
        with torch.no_grad():
            _, _, pred_matte_freeze = modnet_freeze(img, True)
        _, _, pred_matte = modnet(img, False)

        loss = 0
        mask_front = (trimap == 1).float()
        mask_transition = (trimap == 0.5).float()
        mask_bg = (trimap == 0).float()
        # keep batch dimension
        loss_fg = torch.mean(torch.sum((mask_front * F.l1_loss(pred_matte, matte, reduction='none')), dim=(1,2,3)) /\
                    (torch.sum(mask_front, dim=(1,2,3)) + 1e-6))
        loss_transition = torch.mean(torch.sum(mask_transition * F.l1_loss(pred_matte, pred_matte_freeze, reduction='none'), dim=(1,2,3)) /\
                    (torch.sum(mask_transition, dim=(1,2,3))+1e-6))
        loss_bg = torch.mean(torch.sum(mask_bg * pred_matte, dim=(1,2,3))/ (torch.sum(mask_bg, dim=(1,2,3))+1e-6))
        loss = loss_bg + 0.1*loss_transition + loss_fg

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if log_step:
            # Tensorboard image & matte log
            # tb_logger.add_images
            img_plot = torch.cat([(img.cpu()+1)/2, 
                            torch.cat([pred_matte.cpu()]*3, 1),
                            torch.cat([pred_matte_freeze.cpu()]*3, 1),
                            torch.cat([trimap.cpu()]*3, 1)
                            ], 0)
            # assert img_plot.size(0) == 3, img_plot.size()
            tb_logger.add_images('img-pred-pred_freeze-trimap', img_plot, step_image)
            # tb_logger.add_image('image', (image[0]+1)/2, step_image)
            # tb_logger.add_image('fg', pred_fg[0], step_image)
            # tb_logger.add_image('fg_backup', pred_backup_fg[0], step_image)
            step_image += 1
        # Tensorboard loss log
        if (idx+1) % 10 == 0:
            tb_logger.add_scalars('unseen_imageset_loss', {'fg': loss_fg.item(),
                                                            'bg': loss_bg.item(),
                                                            'trans': loss_transition.item()},
                                                step_scalar)
            step_scalar += 1
            
    # 每个Epoch checkpoint一次
    torch.save(modnet.state_dict(), os.path.join(folder_ckpt, 'epoch_%d.pth'%(epoch)))
    
